package com.crazymaker.springcloud.ratelimit;

import lombok.extern.slf4j.Slf4j;
import org.junit.Test;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶算法
 */
@Slf4j
public class TokenBucketLimiter {
    // 上次令牌的发放时间
    private static long lastTime = 0;

    // 令牌生成的速度：每秒2个
    private static final int rate = 2;

    // 时间间隔
    private static final long interval = 1000;

    // 桶的容量
    private static final int capacity = 20;

    // 可用令牌的数量
    private static AtomicInteger tokens = new AtomicInteger(0);

    public synchronized static boolean isLimit(String taskId, int applyCount) {
        long now = System.currentTimeMillis();
        long gap = now - lastTime;

        // 上次距离现在还不到1秒
        // 此1秒内不用更新lastTime
        if (lastTime != 0 && gap < interval) {
            if (tokens.get() < applyCount) {
                return true;
            } else {
                // 放行
                tokens.getAndAdd(-applyCount);
                return false;
            }
        }

        // 超过1秒的范围
        if (lastTime == 0) {
            gap = 1000;
        }

        // 该时间内生成的令牌数
        int times = ((int)gap/1000) * rate;
        int total = tokens.get() + times;
        // 最多只能放行的个数
        tokens.set(Math.max(capacity, total));
        // 更新上次令牌发放时间
        lastTime = now;

        if (tokens.get() < applyCount) {
            return true;
        } else {
            // 放行
            tokens.getAndAdd(-applyCount);
            return false;
        }
    }

    ExecutorService executorService = Executors.newFixedThreadPool(5);

    @Test
    public void testLimit() throws Exception{
        CountDownLatch latch = new CountDownLatch(5);
        AtomicInteger failedCount = new AtomicInteger();
        long start = System.currentTimeMillis();
        for (int i = 0; i < 5; i++) {
            executorService.submit(new Runnable() {
                @Override
                public void run() {
                    try {
                        String taskId = String.valueOf(Thread.currentThread().getId());
                        for (int j = 0; j < 20; j++) {
                            // 此处可以动态调整，获取多少个令牌
                            boolean res = TokenBucketLimiter.isLimit("taskId" + taskId, 2);
                            if (res) {
                                // 失败的次数
                                failedCount.incrementAndGet();
                            }
                            Thread.sleep(200);
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        latch.countDown();
                    }
                }
            });
        }

        latch.await();
        long end = System.currentTimeMillis() - start;
        log.info("失败的次数:{}", failedCount.get());
        log.info("成功的次数:{}", (5 * 20) - failedCount.get());
        log.info("失败的比例:{}", failedCount.get() / 100F);
        log.info("耗时:{}", end / 1000F);
    }
}
